#ifndef _DatasetIO_dataset_io_kitti_H_
#define _DatasetIO_dataset_io_kitti_H_

#include "dataset_io.h"

namespace DatasetIO {

class KittiDataset : public Dataset {
    std::string path;

    size_t num_cams;
    std::vector<double> image_timestamps;
    std::unordered_map<double, std::string> image_path;
    size_t cur_image_id = 0;

    std::vector<AccelData::Ptr> accel_data;
    std::vector<GyroData::Ptr> gyro_data;

    std::vector<double> gt_timestamps;        // ordered gt timestamps
    std::vector<PoseData::Ptr> gt_pose_data;  // TODO: change to eigen aligned

    int64_t mocap_to_imu_offset_ns;

    const std::vector<std::string> folder = {"/image_0/", "/image_1/"};

public:
    KittiDataset() : Dataset() { type = DatasetType::kKitti; }
    ~KittiDataset(){};

    size_t get_num_cams() const { return num_cams; }

    std::vector<double> &get_image_timestamps() { return image_timestamps; }

    const std::vector<AccelData::Ptr> &get_accel_data() const {
        return accel_data;
    }

    const std::vector<GyroData::Ptr> &get_gyro_data() const {
        return gyro_data;
    }

    const std::vector<double> &get_gt_timestamps() const {
        return gt_timestamps;
    }

    const std::vector<PoseData::Ptr> &get_gt_pose_data() const {
        return gt_pose_data;
    }

    int64_t get_mocap_to_imu_offset_ns() const {
        return mocap_to_imu_offset_ns;
    }

    std::vector<MonoImageData::Ptr> get_image_data(double t_ns) {
        std::vector<MonoImageData::Ptr> res(num_cams);

        for (size_t i = 0; i < num_cams; i++) {
            std::string full_image_path = path + folder[i] + image_path[t_ns];

            if (fs::exists(full_image_path)) {
                cv::Mat img;
                if (!(img = cv::imread(full_image_path, cv::IMREAD_UNCHANGED))
                         .empty()) {
                    res[i] = std::make_shared<MonoImageData>();
                    res[i]->timestamp = t_ns;
                    res[i]->data = img.clone();
                }
            }
        }

        return res;
    }

    MonoImageData::Ptr get_next_mono_image_data() {
        MonoImageData::Ptr res;
        cv::Mat image;

        double t_s = image_timestamps[cur_image_id++];

        std::string full_image_path = path + folder[0] + image_path[t_s];
        if (!(image = cv::imread(full_image_path, cv::IMREAD_UNCHANGED))
                 .empty()) {
            res = std::make_shared<MonoImageData>();
            res->timestamp = t_s;
            res->data = image.clone();
        }

        return res;
    }

    EIGEN_MAKE_ALIGNED_OPERATOR_NEW

    friend class KittiIO;
};

class KittiIO : public DatasetIoInterface {
public:
    KittiIO() {}

    void read(const std::string &path) {
        if (!fs::exists(path))
            LOG(WARNING) << "No dataset found in " << path;

        data.reset(new KittiDataset);

        data->num_cams = 2;
        data->path = path;

        read_image_timestamps(path + "/times.txt");

        if (fs::exists(path + "/poses.txt")) {
            read_gt_data_pose(path + "/poses.txt");
        }
    }

    void reset() { data.reset(); }

    Dataset::Ptr get_data() { return data; }

private:
    void read_image_timestamps(const std::string &path) {
        std::ifstream f(path);
        std::string line;
        while (std::getline(f, line)) {
            if (line[0] == '#')
                continue;
            std::stringstream ss(line);

            double t_s;
            ss >> t_s;

            std::stringstream ss1;
            ss1 << std::setfill('0') << std::setw(6)
                << data->image_timestamps.size() << ".png";

            data->image_timestamps.emplace_back(t_s);
            data->image_path[t_s] = ss1.str();
        }
    }

    void read_gt_data_pose(const std::string &path) {
        data->gt_timestamps.clear();
        data->gt_pose_data.clear();

        int i = 0;

        std::ifstream f(path);
        std::string line;
        while (std::getline(f, line)) {
            if (line[0] == '#')
                continue;

            std::stringstream ss(line);

            Eigen::Matrix3d rot;
            Eigen::Vector3d pos;

            ss >> rot(0, 0) >> rot(0, 1) >> rot(0, 2) >> pos[0] >> rot(1, 0) >>
                rot(1, 1) >> rot(1, 2) >> pos[1] >> rot(2, 0) >> rot(2, 1) >>
                rot(2, 2) >> pos[2];

            Eigen::Matrix4d pose = Eigen::Matrix4d::Identity();
            pose.block<3, 3>(0, 0) = rot;
            pose.block<3, 1>(0, 3) = pos;

            PoseData::Ptr pose_data(new PoseData);
            pose_data->timestamp = data->image_timestamps[i];
            pose_data->data = pose;

            data->gt_timestamps.emplace_back(data->image_timestamps[i]);
            data->gt_pose_data.emplace_back(pose_data);
            i++;
        }
    }

    std::shared_ptr<KittiDataset> data;
};

}  // namespace DatasetIO

#endif